import os
import pandas as pd
from pathlib import Path
from colorama import Fore, Back, Style, init
import argparse

def find_and_check_excel_files(specified_file=None):
    # 初始化colorama，确保在Windows系统上也能正常工作
    init(autoreset=True)
    
    # 存储所有发现精度不通过的文件路径
    failed_files = []
    
    # 获取当前脚本所在目录
    script_dir = Path(__file__).parent.resolve()
    # 构建atk_output文件夹路径
    atk_output_dir = os.path.join(script_dir, "atk_output")
    
    # 如果指定了文件，则只处理该文件
    if specified_file:
        # 检查指定文件是否存在且是xlsx文件
        if not os.path.exists(specified_file) or not specified_file.endswith(".xlsx"):
            print(f"{Fore.RED}错误：指定的文件 {specified_file} 不存在或不是xlsx文件")
            return failed_files
            
        print(f"{Fore.CYAN}正在处理指定文件: {specified_file}")
        process_excel_file(specified_file, failed_files)
        return failed_files
    
    # 检查atk_output文件夹是否存在
    if not os.path.exists(atk_output_dir) or not os.path.isdir(atk_output_dir):
        print(f"{Fore.RED}错误：{atk_output_dir} 文件夹不存在或不是一个目录")
        return failed_files
    
    # 遍历atk_output下的所有文件夹
    for root, dirs, files in os.walk(atk_output_dir):
        # 检查当前目录是否包含report子文件夹
        report_dir = os.path.join(root, "report")
        if os.path.exists(report_dir) and os.path.isdir(report_dir):
            # 在report文件夹中查找xlsx文件
            for file in os.listdir(report_dir):
                if file.endswith(".xlsx"):
                    excel_path = os.path.join(report_dir, file)
                    print(f"{Fore.CYAN}正在处理文件: {excel_path}")
                    process_excel_file(excel_path, failed_files)
    
    return failed_files

def process_excel_file(excel_path, failed_files):
    """处理单个Excel文件并检查精度"""
    try:
        # 读取Excel文件
        df = pd.read_excel(excel_path)
        
        # 检查是否包含"cpu_0_精度通过"列
        if "cpu_0_精度通过" not in df.columns:
            print(f"{Fore.YELLOW}  警告：文件中不包含'cpu_0_精度通过'列")
            return
        
        # 筛选出"cpu_0_精度通过"为False的行
        failed_rows = df[df["cpu_0_精度通过"] != True]
        
        if not failed_rows.empty:
            # 将文件添加到失败列表
            failed_files.append(excel_path)
            print(f"{Fore.RED}  发现 {len(failed_rows)} 条CPU精度未通过的记录：")
            # 遍历并输出结果
            for _, row in failed_rows.iterrows():
                # 确保所需列存在
                number = row.get("编号", "未知")
                precision_detail = row.get("精度详情", "无信息")
                json_info = row.get("用例json信息", "无信息")
                
                print(f"    编号: {number}")
                print(f"    精度详情: {precision_detail}")
                print(f"    用例json信息: {json_info}")
                print("    " + "-"*50)
        else:
            # 没有未通过记录时，用绿色背景和黑色文字突出显示
            print(f"{Back.GREEN}{Fore.BLACK}  该文件中所有CPU精度检查均通过！没有未通过记录。{Style.RESET_ALL}")
            
    except Exception as e:
        print(f"{Fore.RED}  处理文件时出错: {str(e)}")

if __name__ == "__main__":
    # 设置命令行参数解析器
    parser = argparse.ArgumentParser(description='检查Excel文件中的CPU精度未通过记录')
    parser.add_argument('-o', '--output', help='指定要查看的单个Excel文件路径')
    args = parser.parse_args()
    
    # 设置pandas显示选项，确保中文正常显示
    pd.set_option('display.unicode.ambiguous_as_wide', True)
    pd.set_option('display.unicode.east_asian_width', True)
    
    # 执行检查并获取失败文件列表
    failed_files = find_and_check_excel_files(args.output)
    
    print(f"\n{Fore.GREEN}处理完成")
    
    # 打印所有精度不通过的文件
    if failed_files:
        print(f"\n{Fore.RED}以下文件存在精度不通过的记录：")
        for i, file in enumerate(failed_files, 1):
            print(f"  {i}. {file}")
    else:
        print(f"\n{Fore.GREEN}所有检查的文件中均未发现精度不通过的记录！")
